#include <torch/extension.h>

/*
template <typename scalar_t>
__global__ void my_cuda_kernel(scalar_t*);
*/
template <typename scalar_t>
__global__ void my_cuda_kernel(torch::PackedTensorAccessor32<scalar_t, 2, torch::RestrictPtrTraits>,
                               torch::PackedTensorAccessor32<scalar_t, 2, torch::RestrictPtrTraits>);

torch::Tensor my_cuda_test(torch::Tensor in_tensor_2)
{
    const auto dim0 = in_tensor_2.size(0);
    const auto dim1 = in_tensor_2.size(1);

    // 2-dimensional thread grid & block
    const dim3 my_blk(128, dim1 / 128);
    const dim3 my_grid(4, dim0 / 4);

    /*
    AT_DISPATCH_FLOATING_TYPES(in_tensor_2.type(), "CUDA exercise", 
                               ([&] {
    my_cuda_kernel<scalar_t><<<my_grid, my_blk>>>(in_tensor_2.data<scalar_t>());
}));
    */
    auto out_tensor_2 = torch::zeros_like(in_tensor_2);
    AT_DISPATCH_FLOATING_TYPES(in_tensor_2.type(), "CUDA exercise", (
        [&] {
                my_cuda_kernel<scalar_t><<<my_grid, my_blk>>>(
                    in_tensor_2.packed_accessor32<scalar_t, 2, torch::RestrictPtrTraits>(),
                    out_tensor_2.packed_accessor32<scalar_t, 2, torch::RestrictPtrTraits>());
        }));
    return out_tensor_2;
    // cudaDeviceSynchronize();
}

/*
template <typename scalar_t>
__global__ void my_cuda_kernel(scalar_t* __restrict__ in_tensor)
*/
template <typename scalar_t>
__global__ void my_cuda_kernel(torch::PackedTensorAccessor32<scalar_t, 2, torch::RestrictPtrTraits> in_tensor, 
                               torch::PackedTensorAccessor32<scalar_t, 2, torch::RestrictPtrTraits> out_tensor)
{
    const int thd_local_idx = threadIdx.x + threadIdx.y * blockDim.x;
    // const int thd_per_blk = blockDim.y * blockDim.x;
    const int blk_idx = blockIdx.x + blockIdx.y * gridDim.x;
    /*
    __shared__ scalar_t shared_mem[sizeof(scalar_t) * 256];

    const scalar_t global_val = in_tensor[blk_idx * thd_per_blk + thd_local_idx];
    shared_mem[thd_local_idx] = global_val;
    __syncthreads();

    // if ((blockIdx.y == 0) && (blockIdx.x == 0) && (threadIdx.y == 0) && (threadIdx.x == 0))
    printf("The number of `scalar_t`'s bytes is %zu!", sizeof(scalar_t));
    */

    scalar_t local_val = in_tensor[blk_idx][thd_local_idx];
    // scalar_t res_val = sqrtf(local_val);
    scalar_t res_val = powf(local_val, 2);
    out_tensor[blk_idx][thd_local_idx] = res_val;
}
